/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.buffer;

import com.aduib.boot.common.log.LOG;

/**
 * 仿照Netty的思路 Chunk由Page组成，是一块连续内存，由memoryMap和depthMap定义成一种平衡二叉树的管理结构
 *
 * @description: BufferChunk
 * @author: zzh
 * @date: 2021/12/1 14:55
 */
public final class BufferChunk {

  private static final int SIZE_BIT_LENGTH = 15;
  private static final int INUSED_BIT_LENGTH = 1;
  private static final int SUBPAGE_BIT_LENGTH = 1;
  private static final int BITMAP_IDX_BIT_LENGTH = 32;

  static final int IS_SUBPAGE_SHIFT = BITMAP_IDX_BIT_LENGTH;
  static final int IS_USED_SHIFT = SUBPAGE_BIT_LENGTH + IS_SUBPAGE_SHIFT;
  static final int SIZE_SHIFT = INUSED_BIT_LENGTH + IS_USED_SHIFT;
  static final int RUN_OFFSET_SHIFT = SIZE_BIT_LENGTH + SIZE_SHIFT;
  // in bytes
  protected final int chunkSize;
  final long bufAddress;
  final Buffer memory;
  final Buffer base;
  final BufferPool bufferPool;
  private final byte[] memoryMap;
  private final byte[] depthMap;
  // in bytes
  private final int pageSize;
  private final int pageShifts;
  private final int chunkPageSize;
  private final int maxOrder;
  private final byte unusable;
  private final int log2PageSize;
  /** manage all subpages in this chunk */
  private final BufferSubpage[] subpages;
  /** store the first page and last page of each avail run */
  private final LongLongHashMap runsAvailMap;
  /** manage all avail runs */
  private final LongPriorityQueue[] runsAvail;
  BufferChunk prev;
  BufferChunk next;
  BufferChunkList parent;
  int freeBytes;

  public BufferChunk(int pageSize, int chunkSize, BufferPool bufferPool, Buffer buffer) {
    this(pageSize, 0, 0, chunkSize, bufferPool, buffer, buffer);
  }

  public BufferChunk(
      int pageSize,
      int pageShifts,
      int maxPageIdx,
      int chunkSize,
      BufferPool bufferPool,
      Buffer memory,
      Buffer base) {
    this.bufferPool = bufferPool;

    this.pageSize = pageSize;
    this.pageShifts = pageShifts;
    this.chunkSize = chunkSize;
    this.chunkPageSize = chunkSize / pageSize;
    this.maxOrder = log2(this.chunkPageSize) + 1;
    this.unusable = (byte) this.maxOrder;
    this.freeBytes = chunkSize;
    this.memory = memory;
    this.base = base;
    this.bufAddress = memory.address();

    this.depthMap = new byte[(1 << this.maxOrder)];
    this.memoryMap = new byte[this.depthMap.length];

    this.log2PageSize = log2(pageSize);

    int memoryMapIndex = 1;
    for (int d = 0; d < maxOrder; ++d) { // move down the tree one level at a time
      int depth = 1 << d;
      for (int p = 0; p < depth; ++p) {
        // in each level traverse left to right and set value to the depth of subtree
        memoryMap[memoryMapIndex] = (byte) d;
        depthMap[memoryMapIndex] = (byte) d;
        memoryMapIndex++;
      }
    }

    runsAvail = newRunsAvailqueueArray(maxPageIdx);
    runsAvailMap = new LongLongHashMap(-1);
    subpages = new BufferSubpage[chunkSize >> pageShifts];

    // insert initial run, offset = 0, pages = chunkSize / pageSize
    int pages = chunkSize >> pageShifts;
    long initHandle = (long) pages << SIZE_SHIFT;
    insertAvailRun(0, pages, initHandle);
  }

  private static LongPriorityQueue[] newRunsAvailqueueArray(int size) {
    LongPriorityQueue[] queueArray = new LongPriorityQueue[size];
    for (int i = 0; i < queueArray.length; i++) {
      queueArray[i] = new LongPriorityQueue();
    }
    return queueArray;
  }

  private static int lastPage(int runOffset, int pages) {
    return runOffset + pages - 1;
  }

  static boolean isUsed(long handle) {
    return (handle >> IS_USED_SHIFT & 1) == 1L;
  }

  private static long toRunHandle(int runOffset, int runPages, int inUsed) {
    return (long) runOffset << RUN_OFFSET_SHIFT
        | (long) runPages << SIZE_SHIFT
        | (long) inUsed << IS_USED_SHIFT;
  }

  static int runOffset(long handle) {
    return (int) (handle >> RUN_OFFSET_SHIFT);
  }

  static int runSize(int pageShifts, long handle) {
    return runPages(handle) << pageShifts;
  }

  static int runPages(long handle) {
    return (int) (handle >> SIZE_SHIFT & 0x7fff);
  }

  static int bitmapIdx(long handle) {
    return (int) handle;
  }

  static boolean isRun(long handle) {
    return !isSubpage(handle);
  }

  static boolean isSubpage(long handle) {
    return (handle >> IS_SUBPAGE_SHIFT & 1) == 1L;
  }

  private static int log2(int chunkSize) {
    if (chunkSize <= 0) {
      LOG.w("invalid parameter!");
      throw new IllegalArgumentException();
    }
    return Integer.SIZE - 1 - Integer.numberOfLeadingZeros(chunkSize);
  }

  private void insertAvailRun(int runOffset, int pages, long handle) {
    int pageIdxFloor = ((BufferPoolArena) bufferPool).pages2pageIdxFloor(pages);
    LongPriorityQueue queue = runsAvail[pageIdxFloor];
    queue.offer(handle);

    // insert first page of run
    insertAvailRun0(runOffset, handle);
    if (pages > 1) {
      // insert last page of run
      insertAvailRun0(lastPage(runOffset, pages), handle);
    }
  }

  private void insertAvailRun0(int runOffset, long handle) {
    long pre = runsAvailMap.put(runOffset, handle);
    assert pre == -1;
  }

  private void removeAvailRun(long handle) {
    int pageIdxFloor = ((BufferPoolArena) bufferPool).pages2pageIdxFloor(runPages(handle));
    LongPriorityQueue queue = runsAvail[pageIdxFloor];
    removeAvailRun(queue, handle);
  }

  private void removeAvailRun(LongPriorityQueue queue, long handle) {
    queue.remove(handle);

    int runOffset = runOffset(handle);
    int pages = runPages(handle);
    // remove first page of run
    runsAvailMap.remove(runOffset);
    if (pages > 1) {
      // remove last page of run
      runsAvailMap.remove(lastPage(runOffset, pages));
    }
  }

  private long getAvailRunByOffset(int runOffset) {
    return runsAvailMap.get(runOffset);
  }

  private int calculateRunSize(int sizeIdx) {
    int maxElements = 1 << pageShifts - BufferPoolArena.LOG2_QUANTUM;
    int runSize = 0;
    int nElements;

    final int elemSize = ((BufferPoolArena) bufferPool).sizeIdx2size(sizeIdx);

    // find lowest common multiple of pageSize and elemSize
    do {
      runSize += pageSize;
      nElements = runSize / elemSize;
    } while (nElements < maxElements && runSize != nElements * elemSize);

    while (nElements > maxElements) {
      runSize -= pageSize;
      nElements = runSize / elemSize;
    }

    assert nElements > 0;
    assert runSize <= chunkSize;
    assert runSize >= elemSize;

    return runSize;
  }

  public boolean isInThisChunk(Buffer buffer) {
    long address = buffer.address();
    return (address >= bufAddress) && (address < bufAddress + chunkSize);
  }

  public int usage() {
    final int freeBytes = this.freeBytes;
    if (freeBytes == 0) {
      return 100;
    }

    int freePercentage = (int) (freeBytes * 100L / chunkSize);
    if (freePercentage == 0) {
      return 99;
    }
    return 100 - freePercentage;
  }

  private int runFirstBestFit(int pageIdx) {
    if (freeBytes == chunkSize) {
      return ((BufferPoolArena) bufferPool).nPSizes - 1;
    }
    for (int i = pageIdx; i < ((BufferPoolArena) bufferPool).nPSizes; i++) {
      LongPriorityQueue queue = runsAvail[i];
      if (queue != null && !queue.isEmpty()) {
        return i;
      }
    }
    return -1;
  }

  private long splitLargeRun(long handle, int needPages) {
    assert needPages > 0;

    int totalPages = runPages(handle);
    assert needPages <= totalPages;

    int remPages = totalPages - needPages;

    if (remPages > 0) {
      int runOffset = runOffset(handle);

      // keep track of trailing unused pages for later use
      int availOffset = runOffset + needPages;
      long availRun = toRunHandle(availOffset, remPages, 0);
      insertAvailRun(availOffset, remPages, availRun);

      // not avail
      return toRunHandle(runOffset, needPages, 1);
    }

    // mark it as used
    handle |= 1L << IS_USED_SHIFT;
    return handle;
  }

  private long allocateSubpage(int sizeIdx) {
    // Obtain the head of the PoolSubPage pool that is owned by the PoolArena and synchronize on it.
    // This is need as we may add it back and so alter the linked-list structure.
    BufferPoolArena arena = (BufferPoolArena) this.bufferPool;
    BufferSubpage head = arena.findSubpageHead(sizeIdx);
    synchronized (head) {
      // allocate a new run
      int runSize = calculateRunSize(sizeIdx);
      // runSize must be multiples of pageSize
      long runHandle = allocateRunNew(runSize);
      if (runHandle < 0) {
        return -1;
      }

      int runOffset = runOffset(runHandle);
      assert subpages[runOffset] == null;
      int elemSize = arena.sizeIdx2size(sizeIdx);

      BufferSubpage subpage =
          new BufferSubpage(
              head, this, pageShifts, runOffset, runSize(pageShifts, runHandle), elemSize);

      subpages[runOffset] = subpage;
      return subpage.allocate();
    }
  }

  private long collapseRuns(long handle) {
    return collapseNext(collapsePast(handle));
  }

  private long collapsePast(long handle) {
    for (; ; ) {
      int runOffset = runOffset(handle);
      int runPages = runPages(handle);

      long pastRun = getAvailRunByOffset(runOffset - 1);
      if (pastRun == -1) {
        return handle;
      }

      int pastOffset = runOffset(pastRun);
      int pastPages = runPages(pastRun);

      // is continuous
      if (pastRun != handle && pastOffset + pastPages == runOffset) {
        // remove past run
        removeAvailRun(pastRun);
        handle = toRunHandle(pastOffset, pastPages + runPages, 0);
      } else {
        return handle;
      }
    }
  }

  private long collapseNext(long handle) {
    for (; ; ) {
      int runOffset = runOffset(handle);
      int runPages = runPages(handle);

      long nextRun = getAvailRunByOffset(runOffset + runPages);
      if (nextRun == -1) {
        return handle;
      }

      int nextOffset = runOffset(nextRun);
      int nextPages = runPages(nextRun);

      // is continuous
      if (nextRun != handle && runOffset + runPages == nextOffset) {
        // remove next run
        removeAvailRun(nextRun);
        handle = toRunHandle(runOffset, runPages + nextPages, 0);
      } else {
        return handle;
      }
    }
  }

  private int calculateStart(int id) {
    int count = 0;
    for (int i = 1; i < depthMap.length; i++) {
      if (depthMap[i] < depthMap[id]) {
        continue;
      } else if (depthMap[i] == depthMap[id]) {
        if (i == id) {
          break;
        } else {
          count += runLength(i);
        }
      } else {
        break;
      }
    }
    return count;
  }

  private int runLength(int id) {
    // represents the size in #bytes supported by node 'id' in the tree
    return 1 << log2(chunkSize) - depthMap[id];
  }

  private int allocateNode(int d) {
    int id = 1;
    int initial = -(1 << d); // has last d bits = 0 and rest all = 1
    byte val = memoryMap[id];
    if (val > d) { // unusable
      return -1;
    }

    while (val < d
        || (id & initial) == 0) { // id & initial == 1 << d for all ids at depth d, for < d it is 0
      id <<= 1;
      val = memoryMap[id];
      if (val > d) {
        id ^= 1;
        val = memoryMap[id];
      }
    }
    byte value = memoryMap[id];
    assert value == d && (id & initial) == 1 << d
        : String.format("val = %d, id & initial = %d, d = %d", value, id & initial, d);
    memoryMap[id] = unusable; // mark as unusable
    updateParentsAlloc(id);
    return id;
  }

  private void updateParentsAlloc(int id) {
    while (id > 1) {
      int parentId = id >>> 1;
      byte val1 = memoryMap[id];
      byte val2 = memoryMap[id ^ 1];
      byte val = val1 < val2 ? val1 : val2;
      memoryMap[parentId] = val;
      id = parentId;
    }
  }

  public synchronized void freeByteBuffer(Buffer buffer) {
    long address = buffer.address();
    int relativeAddress = (int) (address - bufAddress);
    int length = buffer.capacity();

    int depth = maxOrder - 1 - log2(length / pageSize);
    int count = 0;
    int i;
    for (i = 0; i < depthMap.length; i++) {
      if (depthMap[i] == depth) {
        if (count == relativeAddress) {
          break;
        }
        count += length;
      }
      if (depthMap[i] > depth) {
        break;
      }
    }
    free(i);
  }

  private void free(int handle) {
    if (memoryMap[handle] != depthMap[handle]) {
      freeBytes += runLength(handle);
      memoryMap[handle] = depthMap[handle];
      updateParentsFree(handle);
    }
  }

  private void updateParentsFree(int id) {
    int logChild = depthMap[id] + 1;
    while (id > 1) {
      int parentId = id >>> 1;
      byte val1 = memoryMap[id];
      byte val2 = memoryMap[id ^ 1];
      logChild -=
          1; // in first iteration equals log, subsequently reduce 1 from logChild as we traverse up

      if (val1 == logChild && val2 == logChild) {
        memoryMap[parentId] = (byte) (logChild - 1);
      } else {
        byte val = val1 < val2 ? val1 : val2;
        memoryMap[parentId] = val;
      }

      id = parentId;
    }
  }

  @Deprecated
  public synchronized Buffer allocateRun(int normCapacity) {
    if (normCapacity > chunkSize) {
      LOG.w("try to acquire a buffer with larger size than chunkSize!");
      return null;
    }
    int d = this.maxOrder - 2 - (log2(normCapacity) - this.log2PageSize);
    if (d > this.maxOrder - 1) {
      d = maxOrder - 1;
    }
    int id = allocateNode(d);
    if (id < 0) {
      return null;
    }
    freeBytes -= runLength(id);

    int start = calculateStart(id);
    int end = start + runLength(id);

    memory.limit(end);
    memory.position(start);

    return memory.slice();
  }

  long allocateRunNew(int runSize) {
    int pages = runSize >> pageShifts;
    int pageIdx = ((BufferPoolArena) bufferPool).pages2pageIdx(pages);

    synchronized (runsAvail) {
      // find first queue which has at least one big enough run
      int queueIdx = runFirstBestFit(pageIdx);
      if (queueIdx == -1) {
        return -1;
      }

      // get run with min offset in this queue
      LongPriorityQueue queue = runsAvail[queueIdx];
      long handle = queue.poll();

      assert handle != LongPriorityQueue.NO_VALUE && !isUsed(handle) : "invalid handle: " + handle;

      removeAvailRun(queue, handle);

      if (handle != -1) {
        handle = splitLargeRun(handle, pages);
      }

      freeBytes -= runSize(pageShifts, handle);
      return handle;
    }
  }

  public boolean allocate(
      PooledBuffer buf, int reqCapacity, int sizeIdx, PoolThreadCache threadCache) {
    BufferPoolArena arena = (BufferPoolArena) this.bufferPool;
    final long handle;
    if (sizeIdx <= arena.smallMaxSizeIdx) {
      // small
      handle = allocateSubpage(sizeIdx);
      if (handle < 0) {
        return false;
      }
      assert isSubpage(handle);
    } else {
      int runSize = arena.sizeIdx2size(sizeIdx);
      handle = allocateRunNew(runSize);
    }
    initBuf(buf, reqCapacity, threadCache, handle);
    return true;
  }

  void free(long handle, int normCapacity, Buffer nioBuffer) {
    if (isSubpage(handle)) {
      BufferPoolArena arena = (BufferPoolArena) bufferPool;
      int sizeIdx = arena.size2SizeIdx(normCapacity);
      BufferSubpage head = arena.findSubpageHead(sizeIdx);

      int sIdx = runOffset(handle);
      BufferSubpage subpage = subpages[sIdx];
      assert subpage != null && subpage.doNotDestroy;

      // Obtain the head of the PoolSubPage pool that is owned by the PoolArena and synchronize on
      // it.
      // This is need as we may add it back and so alter the linked-list structure.
      synchronized (head) {
        if (subpage.free(head, bitmapIdx(handle))) {
          // the subpage is still used, do not free it
          return;
        }
        assert !subpage.doNotDestroy;
        // Null out slot in the array as it was freed and we should not use it anymore.
        subpages[sIdx] = null;
      }
    }

    // start free run
    int pages = runPages(handle);

    synchronized (runsAvail) {
      // collapse continuous runs, successfully collapsed runs
      // will be removed from runsAvail and runsAvailMap
      long finalRun = collapseRuns(handle);

      // set run as not used
      finalRun &= ~(1L << IS_USED_SHIFT);
      // if it is a subpage, set it to run
      finalRun &= ~(1L << IS_SUBPAGE_SHIFT);

      insertAvailRun(runOffset(finalRun), runPages(finalRun), finalRun);
      freeBytes += pages << pageShifts;
    }
  }

  void initBuf(PooledBuffer buf, int reqCapacity, PoolThreadCache threadCache, long handle) {
    if (isRun(handle)) {
      buf.init(
          this,
          base,
          handle,
          runOffset(handle) << pageShifts,
          reqCapacity,
          runSize(pageShifts, handle),
          threadCache);
    } else {
      initBufWithSubpage(buf, base, handle, reqCapacity, threadCache);
    }
  }

  void initBufWithSubpage(
      PooledBuffer buf,
      Buffer nioBuffer,
      long handle,
      int reqCapacity,
      PoolThreadCache threadCache) {
    int runOffset = runOffset(handle);
    int bitmapIdx = bitmapIdx(handle);

    BufferSubpage s = subpages[runOffset];
    assert s.doNotDestroy;
    assert reqCapacity <= s.elemSize;

    int offset = (runOffset << pageShifts) + bitmapIdx * s.elemSize;
    buf.init(this, base, handle, offset, reqCapacity, s.elemSize, threadCache);
  }

  public void destroy() {}
}
